#include "asyncstuntcpsocket.h"

#include "p2p/stun.h"

namespace cricket
{
	static const size_t kMaxPacketSize = 64 * 1024;
	
	typedef uint16 PacketLength;
	static const size_t kPacketLenSize = sizeof(PacketLength);
	static const size_t kPacketLenOffset = 2;
	static const size_t kBufSize = kMaxPacketSize + kStunHeaderSize;
	static const size_t kTurnChannelDataHdrSize = 4;

	inline bool IsStunMessage(uint16 msg_type)
	{
		// The first two bits of a channel data message are 0b01.
		return (msg_type & 0xC000) ? false : true;
	}

	// AsyncStunTCPSocket
	AsyncStunTCPSocket* AsyncStunTCPSocket::Create(base::AsyncSocket* socket, 
		const base::SocketAddress& bind_address, 
		const base::SocketAddress& remote_address)
	{
		return new AsyncStunTCPSocket(AsyncTCPSocketBase::ConnectSocket(
			socket, bind_address, remote_address), false);
	}

	AsyncStunTCPSocket::AsyncStunTCPSocket(base::AsyncSocket* socket, bool listen)
		: base::AsyncTCPSocketBase(socket, listen, kBufSize)
	{
	}

	int AsyncStunTCPSocket::Send(const void* pv, size_t cb)
	{
		if (cb > kBufSize || cb < kPacketLenSize + kPacketLenOffset)
		{
			SetError(EMSGSIZE);
			return -1;
		}

		// If we are blocking on send, then silently drop this packet
		if (!IsOutBufferEmpty())
			return static_cast<int>(cb);

		int pad_bytes;
		size_t expected_pkt_len = GetExpectedLength(pv, cb, &pad_bytes);

		// Accept only complete STUN/ChannelData packets.
		if (cb != expected_pkt_len)
			return -1;

		AppendToOutBuffer(pv, cb);

		ASSERT(pad_bytes < 4);
		char padding[4] = {0};
		AppendToOutBuffer(padding, pad_bytes);

		int res = FlushOutBuffer();
		if (res <= 0)
		{
			// drop packet if we made no progress
			ClearOutBuffer();
			return res;
		}

		// We claim to have sent the whole thing, even if we only sent partial
		return static_cast<int>(cb);
	}

	void AsyncStunTCPSocket::ProcessInput(char* data, size_t* len)
	{
		base::SocketAddress remote_addr(GetRemoteAddress());

		// STUN packet - First 4 bytes. Total header size is 20 bytes.
		// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		// |0 0|     STUN Message Type     |         Message Length        |
		// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

		// TURN ChannelData
		// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		// |         Channel Number        |            Length             |
		// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

		while (true)
		{
			// We need at least 4 bytes to read the STUN or channelData packet length.
			if (*len < kPacketLenOffset + kPacketLenSize)
				return;

			int pad_bytes;
			size_t expect_pkt_len = GetExpectedLength(data, *len, &pad_bytes);
			size_t actual_length = expect_pkt_len + pad_bytes;

			if (*len < actual_length)
				return;

			SignalReadPacket(this, data, expect_pkt_len, remote_addr);

			*len -= actual_length;
			if (*len > 0)
				memmove(data, data + actual_length, *len);
		}
	}

	void AsyncStunTCPSocket::HandleIncomingConnection(base::AsyncSocket* socket)
	{
		SignalNewConnection(this, new AsyncStunTCPSocket(socket, false));
	}

	size_t AsyncStunTCPSocket::GetExpectedLength(const void* data, size_t len, int* pad_bytes)
	{
		*pad_bytes = 0;

		PacketLength pkt_len = base::GetBE16(static_cast<const char*>(data) + kPacketLenOffset);
		
		size_t expected_pkt_len;
		uint16 msg_type = base::GetBE16(data);
		if (IsStunMessage(msg_type))
		{
			// STUN message
			expected_pkt_len = kStunHeaderSize + pkt_len;
		}
		else
		{
			// TURN ChannelData message.
			expected_pkt_len = kTurnChannelDataHdrSize + pkt_len;
			// From RFC 5766 section 11.5
			// Over TCP and TLS-over-TCP, the ChannelData message MUST be padded to
			// a multiple of four bytes in order to ensure the alignment of
			// subsequent messages.  The padding is not reflected in the length
			// field of the ChannelData message, so the actual size of a ChannelData
			// message (including padding) is (4 + Length) rounded up to the nearest
			// multiple of 4.  Over UDP, the padding is not required but MAY be
			// included.
			if (expected_pkt_len % 4)
				*pad_bytes = 4 - (expected_pkt_len % 4);
		}
		return expected_pkt_len;
	}

} // namespace cricket